import numpy as np

from utils.distance_utils import nn_euclidean_distance, nn_cosine_distance


class NearestNeighborDistanceMetric(object):
    def __init__(self, distance_mode, matching_threshold, budget=None):
        if distance_mode == "euclidean":
            self.use_func = nn_euclidean_distance
        elif distance_mode == "cosine":
            self.use_func = nn_cosine_distance
        else:
            raise ValueError("Invalid mode; must be either 'euclidean' or 'cosine'")
        self.matching_threshold = matching_threshold
        self.budget = budget
        self.samples = {}

    def add_data(self, features: np.ndarray, targets: np.ndarray, active_targets: list):
        """Update the distance metric with new data"""
        for feature, target in zip(features, targets):
            self.samples.setdefault(target, []).append(feature)
            if self.budget is not None:
                self.samples[target] = self.samples[target][-self.budget:]
        self.samples = {k: self.samples[k] for k in active_targets}

    def distance(self, features: np.ndarray, targets:list):
        """计算 features 和 targets 之间的距离"""
        cost_matrix = np.zeros((len(targets), len(features)))
        for i, target in enumerate(targets):
            cost_matrix[i, :] = self.use_func(self.samples[target], features)
        return cost_matrix
